#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time    : 2024/6/10 17:30
@Author  : thezehui@gmail.com
@File    : 1.RunnableParallel使用技巧.py
"""
import dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableParallel
from langchain_openai import ChatOpenAI

dotenv.load_dotenv()

# TDDO通过使用RunnableMaps，可以并行执行多个Runnables，且可以将这些Runnables的输出作为一个映射返回。
from langchain.schema.runnable import RunnableMap
model = ChatOpenAI()
outputParser = StrOutputParser()
joke_chain = ChatPromptTemplate.from_template("tell me a joke about {topic}") | model | outputParser
poem_chain = ChatPromptTemplate.from_template("write a 2-line poem about {topic}") | model | outputParser

map_chain = RunnableMap({"joke": joke_chain, "poem": poem_chain,})

response = map_chain.invoke({"topic": "bear"})
print(response)
# {'joke': 'Why did the bear bring a flashlight to the party? \n\nBecause he wanted to be the light of the party- bear-y funny!', 'poem': 'In the forest deep and still,\nThe bear roams, a silent thrill.'}
